package com.migration.model.migration.scripts;

import com.migration.model.migration.AbstractMultiThreadMigrationScript;
import com.migration.model.migration.CustomExecutableMigrationScript;
import com.migration.model.newdb.RiskStep;
import com.migration.model.newdb.RiskStepGroup;
import com.migration.model.olddb.RaStep;
import com.migration.model.olddb.RaStepDetail;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.criteria.CriteriaBuilder;
import jakarta.persistence.criteria.CriteriaQuery;
import jakarta.persistence.criteria.Root;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 步骤信息表(risk_step)迁移脚本
 */
@Component
public class A023_RiskStepMigrationScript extends AbstractMultiThreadMigrationScript implements CustomExecutableMigrationScript {

    @Autowired
    @Qualifier("oldEntityManagerFactory")
    private EntityManagerFactory oldEntityManagerFactory;

    @Autowired
    @Qualifier("newEntityManagerFactory")
    private EntityManagerFactory newEntityManagerFactory;

    /**
     * 构造函数，初始化迁移配置
     */
    public A023_RiskStepMigrationScript() {
        super(
                "risk_step_all",
                "步骤信息数据迁移",
                "多表关联",
                "risk_step",
                23
        );
    }

    @Override
    protected String generateMigrationSql() {
        return "";
    }

    @Override
    public int executeCustomMigration(DataSource dataSource) {
        // 清空目标表
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        jdbcTemplate.execute(getTruncateSql());
        jdbcTemplate.execute("TRUNCATE TABLE risk_step_group;");

        // 执行自定义迁移
        return executeWithEntityProcessingMultiThread(oldEntityManagerFactory, newEntityManagerFactory);
    }

    protected int executeWithEntityProcessingMultiThread(EntityManagerFactory oldEntityManagerFactory, EntityManagerFactory newEntityManagerFactory) {
        EntityManager oldEntityManager = oldEntityManagerFactory.createEntityManager();
        AtomicInteger rowsProcessed = createAtomicCounter();

        try {
            // 获取所有需要的数据
            List<RaStep> allSteps = getAllSteps(oldEntityManager);
            List<RaStepDetail> allStepDetails = getAllStepDetails(oldEntityManager);

            // 过滤掉没有租户ID的步骤
            List<RaStep> validSteps = allSteps.stream()
                    .filter(step -> step.getTenantId() != null)
                    .collect(Collectors.toList());

            // 计算总批次数
            int totalStepBatches = calculateBatchCount(validSteps.size());
            int totalDetailBatches = calculateBatchCount(allStepDetails.size());
            int totalBatches = totalStepBatches + totalDetailBatches;

            // 创建线程池和计数器
            ExecutorService executorService = createExecutorService();
            CountDownLatch latch = createCountDownLatch(totalBatches);

            // 第一阶段：迁移RaStep到RiskStepGroup
            Map<Long, RiskStepGroup> oldToNewStepMap = new HashMap<>();
            DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");

            // 按批次处理步骤组数据
            for (int i = 0; i < totalStepBatches; i++) {
                int fromIndex = i * batchSize;
                int toIndex = Math.min(fromIndex + batchSize, validSteps.size());
                List<RaStep> stepBatch = validSteps.subList(fromIndex, toIndex);

                executorService.submit(() -> {
                    try {
                        EntityManager newEntityManager = newEntityManagerFactory.createEntityManager();
                        try {
                            newEntityManager.getTransaction().begin();
                            int batchCount = processStepGroupBatch(stepBatch, newEntityManager, formatter, oldToNewStepMap);
                            rowsProcessed.addAndGet(batchCount);
                            newEntityManager.getTransaction().commit();
                        } catch (Exception e) {
                            if (newEntityManager.getTransaction().isActive()) {
                                newEntityManager.getTransaction().rollback();
                            }
                            throw e;
                        } finally {
                            newEntityManager.close();
                            latch.countDown();
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                });
            }

            // 等待步骤组处理完成
            waitForPartialCompletion(latch, totalDetailBatches);

            // 同步oldToNewStepMap，确保所有线程完成写入
            Map<Long, RiskStepGroup> finalStepMap;
            synchronized (oldToNewStepMap) {
                finalStepMap = new HashMap<>(oldToNewStepMap);
            }

            // 第二阶段：迁移RaStepDetail到RiskStep
            // 按批次处理步骤详情数据
            for (int i = 0; i < totalDetailBatches; i++) {
                int fromIndex = i * batchSize;
                int toIndex = Math.min(fromIndex + batchSize, allStepDetails.size());
                List<RaStepDetail> detailBatch = allStepDetails.subList(fromIndex, toIndex);

                executorService.submit(() -> {
                    try {
                        EntityManager newEntityManager = newEntityManagerFactory.createEntityManager();
                        try {
                            newEntityManager.getTransaction().begin();
                            int batchCount = processStepBatch(detailBatch, newEntityManager, formatter, finalStepMap);
                            rowsProcessed.addAndGet(batchCount);
                            newEntityManager.getTransaction().commit();
                        } catch (Exception e) {
                            if (newEntityManager.getTransaction().isActive()) {
                                newEntityManager.getTransaction().rollback();
                            }
                            throw e;
                        } finally {
                            newEntityManager.close();
                            latch.countDown();
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                });
            }

            // 等待所有任务完成
            waitForCompletion(latch, executorService);

        } catch (Exception e) {
            throw e;
        } finally {
            oldEntityManager.close();
        }

        return rowsProcessed.get();
    }

    /**
     * 处理步骤组批次数据
     */
    private int processStepGroupBatch(List<RaStep> stepBatch, EntityManager newEntityManager,
                                     DateTimeFormatter formatter, Map<Long, RiskStepGroup> oldToNewStepMap) {
        List<RiskStepGroup> stepGroupEntities = new ArrayList<>();

        for (RaStep oldStep : stepBatch) {
            RiskStepGroup stepGroup = new RiskStepGroup();
            stepGroup.setId(oldStep.getId());
            stepGroup.setTenantId(1L);
            stepGroup.setStepGroupName(oldStep.getStepName());
            stepGroup.setStepGroupDesc(oldStep.getStepDesc());

            // 格式化日期时间
            if (oldStep.getCreateTime() != null) {
                stepGroup.setCreateTime(oldStep.getCreateTime().format(formatter));
            }
            if (oldStep.getUpdateTime() != null) {
                stepGroup.setUpdateTime(oldStep.getUpdateTime().format(formatter));
            }
            stepGroup.setCreateBy(1L); // 默认值
            stepGroup.setUpdateBy(1L); // 默认值
            stepGroup.setIsDel(oldStep.getIsDel());

            newEntityManager.persist(stepGroup);
            stepGroupEntities.add(stepGroup);

            // 保存旧ID和新ID的映射关系
            synchronized (oldToNewStepMap) {
                oldToNewStepMap.put(oldStep.getId(), stepGroup);
            }
        }

        return stepGroupEntities.size();
    }

    /**
     * 处理步骤详情批次数据
     */
    private int processStepBatch(List<RaStepDetail> detailBatch, EntityManager newEntityManager,
                               DateTimeFormatter formatter, Map<Long, RiskStepGroup> stepMap) {
        List<RiskStep> stepEntities = new ArrayList<>();

        for (RaStepDetail oldStepDetail : detailBatch) {
            // 检查是否有对应的Step分组
            RiskStepGroup riskStepGroup = stepMap.get(oldStepDetail.getStepId());
            if (riskStepGroup == null) {
                continue; // 如果没有对应的step分组，跳过此记录
            }

            RiskStep step = new RiskStep();
            step.setId(oldStepDetail.getId());
            step.setTenantId(1L);
            step.setStepGroupId(riskStepGroup.getId().intValue());
            step.setStepName(oldStepDetail.getChildStepName());
            step.setStepDesc(null); // 原始数据中没有对应的字段

            // 使用当前时间作为创建和更新时间
            LocalDateTime now = LocalDateTime.now();
            step.setCreateTime(now.format(formatter));
            step.setUpdateTime(now.format(formatter));

            // 设置默认的创建者和更新者
            step.setCreateBy(1L);
            step.setUpdateBy(1L);

            step.setIsDel(oldStepDetail.getIsDel());

            newEntityManager.persist(step);
            stepEntities.add(step);
        }

        return stepEntities.size();
    }

    /**
     * 获取所有未删除的步骤数据
     *
     * @param entityManager 旧数据库实体管理器
     * @return 步骤数据列表
     */
    private List<RaStep> getAllSteps(EntityManager entityManager) {
        CriteriaBuilder cb = entityManager.getCriteriaBuilder();
        CriteriaQuery<RaStep> cq = cb.createQuery(RaStep.class);
        Root<RaStep> root = cq.from(RaStep.class);
        cq.where(cb.equal(root.get("isDel"), (byte) 0));
        cq.select(root);
        return entityManager.createQuery(cq).getResultList();
    }

    /**
     * 获取所有未删除的步骤详情数据
     *
     * @param entityManager 旧数据库实体管理器
     * @return 步骤详情数据列表
     */
    private List<RaStepDetail> getAllStepDetails(EntityManager entityManager) {
        CriteriaBuilder cb = entityManager.getCriteriaBuilder();
        CriteriaQuery<RaStepDetail> cq = cb.createQuery(RaStepDetail.class);
        Root<RaStepDetail> root = cq.from(RaStepDetail.class);
        cq.where(cb.equal(root.get("isDel"), (byte) 0));
        cq.select(root);
        return entityManager.createQuery(cq).getResultList();
    }
}
